Skip to content

[python] Preserve LeRobot recording metadata in direct writer - #9708

Merged
JingsongLi merged 4 commits into
apache:masterfrom
YannByron:codex/lerobot-writer-companion-tables
Sep 10, 2026
Merged

[python] Preserve LeRobot recording metadata in direct writer#9708
JingsongLi merged 4 commits into
apache:masterfrom
YannByron:codex/lerobot-writer-companion-tables

Conversation

@YannByron

@YannByron YannByron commented Sep 9, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

Follow up on #9665 by preserving native LeRobot recording metadata when frames are captured directly into Paimon. The writer now publishes a managed table group containing frames, episodes, tasks, optional subtasks, dataset information, and statistics without creating an intermediate LeRobot dataset directory.

Changes

  • Create and validate __episodes, __tasks, __info, and __stats companion tables for direct capture.
  • Compute native LeRobot episode statistics when an episode is accepted and incrementally maintain global statistics across commits.
  • Aggregate each feature using its declared dtype, avoiding LeRobot 0.4.4's image-name heuristic for ordinary numeric fields.
  • Store task labels once in the tasks table while frame rows retain task_index.
  • Optionally persist an immutable ordered vocabulary in __subtasks, validate frame-level subtask_index values, and restore the vocabulary on resume.
  • Resume frame, episode, task, subtask, and statistics state from snapshot properties and projected companion metadata without scanning frame data.
  • Commit companion metadata before frames and reject inconsistent or legacy frame-only table groups instead of silently continuing from partial state.
  • Document the direct-capture table group, statistics, subtasks, resume behavior, and multi-table commit boundary.

Testing

  • Run all LeRobot loader and writer unit tests (77 tests).
  • Verify numeric and image episode/global statistics against native LeRobot implementations.
  • Verify a numeric observation.image_embedding feature across multiple episodes and resume.
  • Verify multiple flushes, resume, schema validation, task deduplication, subtask vocabulary persistence, and inconsistent companion state handling.
  • Simulate an environment without the LeRobot extra: dependency-requiring tests skip while the missing-dependency test still runs.
  • Run Ruff on all changed Python files.
  • Run git diff --check.

Notes

  • Direct capture stores image frames as PNG blobs; video remains unsupported.
  • A new table declaring subtask_index requires an ordered, unique subtasks vocabulary. Existing groups restore it from __subtasks and reject conflicting vocabularies.
  • Existing frame-only writer tables are not implicitly migrated because they do not contain the metadata required to reconstruct complete native LeRobot recording state.
  • Paimon does not provide a transaction across the table group, so reopening validates component state and rejects a partial batch.

@YannByron Codex @codex
feat(python): persist LeRobot recording metadata
f730404 
Write episodes, tasks, info, and native LeRobot statistics alongside directly captured frames. Restore writer state from companion metadata without scanning the frame table.

Co-Authored-By: Codex <noreply@anthropic.com>

AI-Model: gpt-5
Co-Authored-By: Codex <noreply@ai-tool.com>
Co-Authored-By: Codex <noreply@openai.com>
AI-Contributed/Feature: 518/518
AI-Contributed/UT: 330/330
JingsongLi
JingsongLi reviewed Sep 10, 2026
View reviewed changes

@JingsongLi JingsongLi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two regressions found in the direct writer changes: optional-dependency handling in the test matrix and aggregation of numeric features whose names contain image.

Comment thread paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py
Comment thread paimon-python/pypaimon/multimodal/lerobot/writer.py Outdated
YannByron and others added 2 commits September 10, 2026 11:29
@YannByron Codex @codex
fix(python): harden LeRobot writer statistics
402e61f 
Keep optional-dependency tests isolated and aggregate statistics by declared feature dtype so numeric fields containing image in their names remain valid.

Co-Authored-By: Codex <noreply@anthropic.com>

AI-Model: gpt-5
Co-Authored-By: Codex <noreply@ai-tool.com>
Co-Authored-By: Codex <noreply@openai.com>
AI-Contributed/Feature: 19/19
AI-Contributed/UT: 45/45
@YannByron Codex @codex
feat(python): support LeRobot writer subtasks
0e41e7b 
Persist an ordered subtask vocabulary in a managed companion table, validate frame-level subtask indices, and restore the immutable mapping when appending to an existing table group.

Co-Authored-By: Codex <noreply@anthropic.com>

AI-Model: gpt-5
Co-Authored-By: Codex <noreply@ai-tool.com>
Co-Authored-By: Codex <noreply@openai.com>
AI-Contributed/Feature: 170/170
AI-Contributed/UT: 129/129
JingsongLi
JingsongLi reviewed Sep 10, 2026
View reviewed changes

@JingsongLi JingsongLi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The previous two findings are fixed. One integration issue remains with the PaimonLeRobotDataset reader now merged into master in #9498.

Comment thread paimon-python/pypaimon/multimodal/lerobot/writer.py Outdated
Comment on lines +381 to +382
"tasks_table": pa.Table.from_pylist(
[], schema=_EMPTY_TASKS_SCHEMA),

@JingsongLi JingsongLi Sep 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Preserve the text index metadata for tasks and subtasks

These companion tables store labels as ordinary columns without pandas index metadata, but the PaimonLeRobotDataset reader merged in #9498 restores a DataFrame and uses its index as the task/subtask text (_component_dataframe and _index_names). Consequently, a recording with task="pick" is decoded as task "0", and its first frame fails _validate_control_row with Paimon task at LeRobot index 0 is not assigned to Episode 0, because the episode correctly lists "pick". Subtask labels are likewise decoded as numeric strings. I reproduced this using the PR writer's persisted output and the current master reader; restoring the text indexes makes the same frame pass validation. Please preserve the native text-index pandas metadata for both companion schemas using the existing _prepare_metadata_tables mechanism, or make the reader explicitly support these plain text columns, and add a writer-to-reader regression test.

@YannByron Codex @codex
fix(python): preserve LeRobot text index metadata
3afd498 
Store task and subtask labels as Pandas indexes so the Paimon LeRobot reader restores their text mappings correctly.\n\nCo-Authored-By: Codex <noreply@anthropic.com>\nAI-Model: gpt-5\nAI-Contributed/Feature: 27/27\nAI-Contributed/UT: 13/13
Co-Authored-By: Codex <noreply@ai-tool.com>
Co-Authored-By: Codex <noreply@openai.com>
AI-Contributed/Feature: 27/27
AI-Contributed/UT: 13/13
JingsongLi
JingsongLi reviewed Sep 10, 2026
View reviewed changes

@JingsongLi JingsongLi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The text-index compatibility issue is fixed. One configuration-dependent regression remains in the new metadata helper.

Comment thread paimon-python/pypaimon/multimodal/lerobot/writer.py
return pa.Table.from_pandas(pd.DataFrame(
{component + "_index": np.asarray(indices, dtype=np.int64)},
index=pd.Index(
labels, dtype="string", name=component,

@JingsongLi JingsongLi Sep 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Pin the string storage used for task and subtask labels

pd.Index(..., dtype="string") follows the global pd.options.mode.string_storage setting. With the supported pandas 2.3.3 / PyArrow 19.0.1 combination and mode.string_storage="pyarrow", both empty and populated label indexes become Arrow large_string, which Paimon's schema parser does not support. Creating a writer then fails with Unsupported pyarrow type: large_string after creating part of the table group, and reopening an existing table fails its companion-schema check. Changing this setting between construction and flush() also causes a schema mismatch and makes the writer terminal. I verified that the previous commit succeeds under the same setting. Please pin the index dtype to pd.StringDtype(storage="python"), or otherwise explicitly produce pa.string() while preserving pandas index metadata. The explicit Python storage dtype passed create/write/resume and Unicode task/subtask label round trips with the global setting still on "pyarrow".

@JingsongLi

JingsongLi commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

+1

@JingsongLi
JingsongLi merged commit 3988f2a into apache:master Sep 10, 2026
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@JingsongLi JingsongLi JingsongLi left review comments

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@YannByron @JingsongLi